Adding different lx layout to micro-benchmarks - #807
Conversation
Signed-off-by: Burkhard Ringlein <ngl@zurich.ibm.com> Co-authored-by: Opus 5 / Kilo
|
👋 Hi! Thank you for contributing. We also recommend installing prek and configuring it to check your code before every local commit. |
|
/claude-review (testing this out for the first time) |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| from spyre_inference.custom_ops.utils import convert | ||
| from spyre_inference.v1.attention.backends.spyre_attn import slot_major_kv_layout | ||
| from spyre_inference.v1.attention.backends.spyre_attn import ( | ||
| head_major_kv_layout, |
There was a problem hiding this comment.
head_major_kv_layout does not exist in spyre_attn.py on this branch (nor anywhere in the repo) — it is introduced by the unmerged #783. This import is unconditional at the top of build_inputs_from_requests, which runs for every config, so it raises ImportError on any invocation, including the pre-existing plain and slot_major baselines that used to work. The whole microbench is broken standalone.
The same dependency applies to _fill_folded_cache: on this branch do_kv_cache_update calls self._reshape_fn(...) with a single tensor, so passing the per-kv-head list from _folded_slot_mapping would also fail without #783.
Either land #783 first (and note the ordering), or guard the import so baseline runs still work.
| Shipped sweep configs for the LX-vs-baseline comparison (each run twice, once with | ||
| `--kv-layout lx`): `lx_study_decode_ctx`, `lx_study_prefill`, `lx_study_decode_batch`, | ||
| `lx_study_num_blocks`. Analysis lives in `spyre_attn_analysis_lx_kv_layout.ipynb`. |
There was a problem hiding this comment.
These four configs and the notebook don't exist in the PR. The only config added is lx_study_chunked_prefill.json, which isn't mentioned here. So anyone following the README hits a missing-file error, and the config that ships is undocumented.
Either add lx_study_decode_ctx/lx_study_prefill/lx_study_decode_batch/lx_study_num_blocks and spyre_attn_analysis_lx_kv_layout.ipynb, or update this section to describe lx_study_chunked_prefill.json.
…anch Folds in the pieces from the parallel lx-kv-microbench branch and lets this land on main ahead of torch-spyre#783. - Make the head_major_kv_layout import lazy. It was unconditional, so every microbench run failed to import on a checkout without torch-spyre#783, not just --kv-layout lx. Keeping it lazy is also what allows a pre-torch-spyre#783 baseline arm. - Add --staging-rows / staging_rows: shrink the impl's staging buffers without changing the measured shape, to split kernel cost from staging cost. - Ship the four remaining sweep configs (context, batch, mixed, pin probe). - README: correct the kernel name to _lx_page_attn_kernel, drop references to configs and a notebook that are not in the tree, and list what is. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
3 participants
Description
To be able to benchmark #783
Related Issues
Test Plan
Checklist
bash format.sh)Signed-off-by:line (DCO compliance)